import cv2 as cv
import numpy as np
import os
import sys
from numpy.core.defchararray import equal

from numpy.lib.npyio import NpzFile

def show_yuv(in_filepath):
    """
    :param in_filepath
    """
    
    name = input("Enter your file name:")
    infile = []
    try:
        infile = open(name,"rb")
    except:
        print("Invalid file")
        breakpoint
    width = input("Enter width:")
    width = int(width)

    height = input("Enter height:")
    height = int(height)

    pix_fmt = input("输入文件格式(支持NV21,NV12,I420,YV12):")
    pix_fmt = pix_fmt.upper()
    support_fmts = [cv.COLOR_YUV2BGR_I420,cv.COLOR_YUV2BGR_YV12,cv.COLOR_YUV2BGR_NV21,cv.COLOR_YUV2BGR_NV12]
    cv_fmt = -1
    if pix_fmt == "NV21":
        cv_fmt = cv.COLOR_YUV2BGR_NV21
    elif pix_fmt == "NV12":
        cv_fmt = cv.COLOR_YUV2BGR_NV12
    elif pix_fmt == "I420":
        cv_fmt = cv.COLOR_YUV2BGR_I420
    elif pix_fmt == "YV12":
        cv_fmt = cv.COLOR_YUV2BGR_YV12
    else:
        print("Un-support pixel format!!!")
    if cv_fmt not in support_fmts:
        print("退出程序")
        return    

    data_size = width * height * 3//2;

    data = np.zeros((data_size,),dtype='uint8',order='C')

    for i in range(data_size):
        data[i] = ord(infile.read(1))
    
    infile.close()

    data = np.reshape(data,(height * 3//2,width)).astype('uint8')    
    img = cv.cvtColor(data,cv_fmt)

    cv.imshow("img",img)
    cv.waitKey(0)
    print("end")

    

    # print("目标:"+dst_pix_fmt.upper())
    # print("w,h = {}x{}".format(width,height))

show_yuv("new.1080x1436.nv21")
    






